#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/Dialect/X86Vector/X86VectorDialect.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/Index/IR/IndexDialect.h"
#include "mlir/Dialect/Index/IR/IndexOps.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/IR/Attributes.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Pass/PassRegistry.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "mlir/Transforms/DialectConversion.h"
#include "mlir/Conversion/LLVMCommon/Pattern.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"

#include "llvm/ADT/STLExtras.h"
#include "llvm/Support/Error.h"
#include "llvm/Support/FormatVariadic.h"

#include "crypto/Crypto.h"
#include "sst/SST.h"
#include "core/JsonHardwareInfo.hpp"

namespace mlir {
#include "mlir/Conversion/Passes.h.inc"
} // namespace mlir


using namespace mlir;


//===----------------------------------------------------------------------===//
// Operation conversion
//===----------------------------------------------------------------------===//

namespace {

class Sha1LoweringSSTPattern : public OpRewritePattern<crypto::Sha1multi_OP> {
public:
  using OpRewritePattern<crypto::Sha1multi_OP>::OpRewritePattern;

  LogicalResult matchAndRewrite(crypto::Sha1multi_OP op,
                                PatternRewriter &rewriter) const override {
    auto loc = op->getLoc();

    // get values needed from for-loop body for launching GPU kernel
    Value plaintext = op->getOperand(0);
    Value plaintext_length = op->getOperand(1);
    Value ciphertext = op->getOperand(2);
    Value num = op->getOperand(3);

    Value c0 = rewriter.create<arith::ConstantIndexOp>(loc, 0);
    Value c1 = rewriter.create<arith::ConstantIndexOp>(loc, 1);
    Value c8 = rewriter.create<arith::ConstantIndexOp>(loc, 8);
    Value c20 = rewriter.create<arith::ConstantIndexOp>(loc, 20);
    Value c48 = rewriter.create<arith::ConstantIndexOp>(loc, 48);
    Value c64 = rewriter.create<arith::ConstantIndexOp>(loc, 64);
    Value c88 = rewriter.create<arith::ConstantIndexOp>(loc, 88);
    Value c256 = rewriter.create<arith::ConstantIndexOp>(loc, 256);
    Value cn = rewriter.create<arith::IndexCastOp>(loc, /*out=*/rewriter.getIndexType(), /*in=*/plaintext_length);
    Value cN = rewriter.create<arith::IndexCastOp>(loc, /*out=*/rewriter.getIndexType(), /*in=*/num);
    // Type used
    auto dynI8MemrefType = MemRefType::get({mlir::ShapedType::kDynamic}, rewriter.getI8Type());

    // register fatbinary
    Value handle = rewriter.create<sst::RegisterFatbinOp>(loc, /*type=*/rewriter.getIndexType());
    rewriter.create<sst::RegisterFuncOp>(loc, /*handle=*/handle, /*hostFunc=*/c1);

    // Task split for multiple CPUs
    wafer::Root& HardwareInfo = wafer::Root::getInstance();
    Value cpu_num;
    std::vector<float> vRatioPartialSum;
    if (HardwareInfo.isInitialized()) {
      float sum = 0.0;
      cpu_num = rewriter.create<arith::ConstantIndexOp>(loc, HardwareInfo.getArchitecture().getChips().CPUNum());
      for (wafer::Chip c :HardwareInfo.getArchitecture().getChips().getChipArray()) {
        if (c.getType().has_value() && (*c.getType()).compare("CPU") == 0 ) {
          vRatioPartialSum.push_back(sum);
          sum += c.getDutyRatio();
        }
      }
    }
    else {
      printf("Using single CPU");
      vRatioPartialSum.push_back(0.0);
      cpu_num = c1;
    }
    vRatioPartialSum.push_back(1.0);
    llvm::ArrayRef<float> valuesRef(vRatioPartialSum);
    auto vectorType = VectorType::get({static_cast<int64_t>(vRatioPartialSum.size())}, rewriter.getF32Type());
    DenseElementsAttr arrayAttr = DenseElementsAttr::get(vectorType, valuesRef);
    Value arrayConstant = rewriter.create<arith::ConstantOp>(loc, arrayAttr.getType(), arrayAttr);

    Value cpu_id = rewriter.create<sst::GetCpuIdOp>(loc, /*cpuId=*/rewriter.getIndexType());
    Value lowerRatio = rewriter.create<vector::ExtractElementOp>(loc, arrayConstant, cpu_id);
    Value upperRatio = rewriter.create<vector::ExtractElementOp>(loc, arrayConstant, rewriter.create<arith::AddIOp>(loc, cpu_id, c1));

    Value Nfp = rewriter.create<arith::IndexCastOp>(loc, rewriter.getI32Type(), cN);
    Nfp = rewriter.create<arith::SIToFPOp>(loc, rewriter.getF32Type(), Nfp);
    Value lowerF = rewriter.create<arith::MulFOp>(loc, Nfp, lowerRatio);
    Value lowerI = rewriter.create<arith::FPToSIOp>(loc, rewriter.getI32Type(), lowerF);
    Value lowerBound = rewriter.create<arith::IndexCastOp>(loc, rewriter.getIndexType(), lowerI);
  
    Value upperF = rewriter.create<arith::MulFOp>(loc, Nfp, upperRatio);
    Value upperI = rewriter.create<arith::FPToSIOp>(loc, rewriter.getI32Type(), upperF);
    Value upperBound = rewriter.create<arith::IndexCastOp>(loc, rewriter.getIndexType(), upperI);

    Value batch_size = rewriter.create<arith::SubIOp>(loc, /*lhs=*/upperBound, /*rhs=*/lowerBound);

    // sst allocation
    Value msg_size = rewriter.create<arith::MulIOp>(loc, /*lhs=*/cn, /*rhs=*/batch_size);
    Value hash_size = rewriter.create<arith::MulIOp>(loc, /*lhs=*/batch_size, /*rhs=*/c20);
    Value d_msg = rewriter.create<sst::MallocOp>(loc, /*res=*/dynI8MemrefType, /*size=*/msg_size);
    Value d_hash = rewriter.create<sst::MallocOp>(loc, /*res=*/dynI8MemrefType, /*size=*/hash_size);

    // sst memcopy from host to device (flag 0 means host to device)
    Value h_msg_part = rewriter.create<memref::ViewOp>(loc, /*resultType0=*/dynI8MemrefType, /*source=*/plaintext, /*byte_shift=*/rewriter.create<arith::MulIOp>(loc, /*lhs=*/lowerBound, /*rhs=*/cn), /*sizes=*/ValueRange{msg_size});
    rewriter.create<sst::MemcpyOp>(loc, /*dst=*/d_msg, /*src=*/h_msg_part, /*size=*/msg_size, /*dir=*/rewriter.create<arith::ConstantOp>(loc, rewriter.getI1Type(), rewriter.getIntegerAttr(rewriter.getI1Type(), APInt(1, 0))));

    // sst configure kernel
    rewriter.create<sst::ConfigureCallOp>(loc, /*bx=*/c64, /*by=*/c1, /*bz=*/c1, /*gx=*/c1, /*gy=*/c1, /*gz=*/c1, /*sharedMem=*/c0);

    // sst kernel function argument
    rewriter.create<sst::SetupArgumentOp>(loc, /*arg=*/cn, /*offset=*/c0);
    rewriter.create<sst::SetupArgumentOp>(loc, /*arg=*/d_msg, /*offset=*/c8);
    rewriter.create<sst::SetupArgumentOp>(loc, /*arg=*/d_hash, /*offset=*/c48);
    rewriter.create<sst::SetupArgumentOp>(loc, /*arg=*/batch_size, /*offset=*/c88);

rewriter.create<vector::PrintOp>(loc, cpu_id);
#ifdef _WAFER_DEBUG_
    rewriter.create<vector::PrintOp>(loc, cpu_id);
    rewriter.create<vector::PrintOp>(loc, lowerBound);
    rewriter.create<vector::PrintOp>(loc, batch_size);
#endif

    // sst launch kernel
    rewriter.create<sst::LaunchOp>(loc, /*funcId=*/c1);

    // sst memcopy from device to host (flag 1 means device to host)
    Value h_hash_part = rewriter.create<memref::ViewOp>(loc, /*resultType0=*/dynI8MemrefType, /*source=*/ciphertext, /*byte_shift=*/rewriter.create<arith::MulIOp>(loc, /*lhs=*/lowerBound, /*rhs=*/c20), /*sizes=*/ValueRange{hash_size});
    rewriter.create<sst::MemcpyOp>(loc, /*dst=*/h_hash_part, /*src=*/d_hash, /*size=*/hash_size, /*dir=*/rewriter.create<arith::ConstantOp>(loc, rewriter.getI1Type(), rewriter.getIntegerAttr(rewriter.getI1Type(), APInt(1, 1))));

    rewriter.eraseOp(op);
    return success();
  }
};

} // namespace


void populateConvertCryptoToSSTPatterns(RewritePatternSet &patterns) {
  patterns.add<Sha1LoweringSSTPattern>(patterns.getContext());
}

namespace {
class LowerCryptoToSSTPass : public PassWrapper<LowerCryptoToSSTPass, OperationPass<ModuleOp>> {
public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LowerCryptoToSSTPass)
  LowerCryptoToSSTPass() = default;
  LowerCryptoToSSTPass(const LowerCryptoToSSTPass &) {}

  StringRef getArgument() const final { return "lower-crypto-to-sst"; }
  StringRef getDescription() const final { return "Lower Crypto Dialect to SST."; }

  void runOnOperation() override;

  void getDependentDialects(DialectRegistry &registry) const override {
    registry
        .insert<func::FuncDialect,
                mlir::LLVM::LLVMDialect,
                memref::MemRefDialect,
                vector::VectorDialect,
                arith::ArithDialect,
                affine::AffineDialect,
                scf::SCFDialect,
                index::IndexDialect,
                sst::SSTDialect
                >();
  }
};
} // end anonymous namespace.

void LowerCryptoToSSTPass::runOnOperation() {
  MLIRContext *context = &getContext();
  ModuleOp module = getOperation();

  ConversionTarget target(*context);
  target.addLegalDialect<
  func::FuncDialect,
  vector::VectorDialect,
  mlir::LLVM::LLVMDialect,
  memref::MemRefDialect,
  arith::ArithDialect,
  affine::AffineDialect,
  scf::SCFDialect,
  index::IndexDialect,
  sst::SSTDialect
  >();
  target.addLegalOp<
  ModuleOp,
  func::FuncOp,
  func::ReturnOp
  >();

  RewritePatternSet patterns(context);
  populateConvertCryptoToSSTPatterns(patterns);

  if (failed(applyPartialConversion(module, target, std::move(patterns))))
    signalPassFailure();
}

void registerLowerCryptoToSSTPass() { PassRegistration<LowerCryptoToSSTPass>(); }

std::unique_ptr<Pass> createLowerCryptoToSSTPass(){ return std::make_unique<LowerCryptoToSSTPass>(); }
